import os
import shutil
import threading
from pathlib import Path

import cv2
import fitz  # PyMuPDF
import numpy as np
from PIL import Image

from app.config import settings
from app.llm import describe_image_with_llm


def _has_colored_frame(
        image_cv, color_name, border_thickness=10, threshold_ratio=0.5
):
    """
    检查图像是否有指定颜色的边框。

    Args:
        image_cv: OpenCV格式的图像 (BGR)。
        color_name: 颜色名称 ('red', 'green', 等)。
        border_thickness: 要检查的边框厚度（像素）。
        threshold_ratio: 边框区域中颜色像素必须达到的比例才算成功。

    Returns:
        bool: 如果存在指定颜色的边框，则为 True。
    """
    h, w, _ = image_cv.shape
    if h < border_thickness * 2 or w < border_thickness * 2:
        return False

    # 将图像从 BGR 转换为 HSV
    hsv_image = cv2.cvtColor(image_cv, cv2.COLOR_BGR2HSV)

    # 定义不同颜色的 HSV 范围
    # 注意：OpenCV中 H范围是 [0, 179], S是 [0, 255], V是 [0, 255]
    colors = {
        "red": ([0, 70, 50], [10, 255, 255]),  # 低色调范围
        "red2": ([170, 70, 50], [180, 255, 255]),  # 高色调范围 (红色环绕)
        "green": ([35, 70, 50], [85, 255, 255]),
        "blue": ([100, 70, 50], [130, 255, 255]),
    }

    if color_name not in colors:
        return False

    # 创建颜色掩码
    lower, upper = colors[color_name]
    mask = cv2.inRange(hsv_image, np.array(lower), np.array(upper))

    # 如果是红色，需要合并两个范围
    if color_name == "red":
        lower2, upper2 = colors["red2"]
        mask2 = cv2.inRange(hsv_image, np.array(lower2), np.array(upper2))
        mask = cv2.bitwise_or(mask, mask2)

    # 检查四个边框区域
    # 1. 上边框
    top_border = mask[0:border_thickness, :]
    # 2. 下边框
    bottom_border = mask[h - border_thickness: h, :]
    # 3. 左边框
    left_border = mask[:, 0:border_thickness]
    # 4. 右边框
    right_border = mask[:, w - border_thickness: w]

    # 计算每个边框中颜色像素的比例
    top_ratio = cv2.countNonZero(top_border) / top_border.size
    bottom_ratio = cv2.countNonZero(bottom_border) / bottom_border.size
    left_ratio = cv2.countNonZero(left_border) / left_border.size
    right_ratio = cv2.countNonZero(right_border) / right_border.size

    # 如果所有边框的颜色像素比例都超过阈值，则认为是带框的
    if all(
            ratio > threshold_ratio
            for ratio in [top_ratio, bottom_ratio, left_ratio, right_ratio]
    ):
        return True

    return False


def _merge_overlapping_regions(regions):
    """合并重叠的区域"""
    if not regions:
        return []

    # 按面积排序，保留较大的区域
    regions = sorted(regions, key=lambda r: (r[2] - r[0]) * (r[3] - r[1]), reverse=True)
    merged = []

    for region in regions:
        x1, y1, x2, y2 = region
        overlapped = False

        for i, existing in enumerate(merged):
            ex1, ey1, ex2, ey2 = existing

            # 检查重叠
            if not (x2 < ex1 or x1 > ex2 or y2 < ey1 or y1 > ey2):
                # 合并区域
                merged[i] = (min(x1, ex1), min(y1, ey1), max(x2, ex2), max(y2, ey2))
                overlapped = True
                break

        if not overlapped:
            merged.append(region)

    return merged


def extract_text_blocks(page):
    """提取文本块，用于区分文本和图表区域"""
    blocks = page.get_text("dict")
    text_regions = []

    for block in blocks["blocks"]:
        if "lines" in block:  # 文本块
            bbox = block["bbox"]
            text_regions.append(bbox)

    return text_regions


def is_likely_diagram(region, page_image, text_regions):
    """
    判断区域是否可能是架构图

    Args:
        region: (x1, y1, x2, y2) 区域坐标
        page_image: 页面图像
        text_regions: 文本区域列表
    """
    x1, y1, x2, y2 = region

    # 提取区域图像
    roi = page_image.crop((x1, y1, x2, y2))
    roi_cv = cv2.cvtColor(np.array(roi), cv2.COLOR_RGB2BGR)
    gray_roi = cv2.cvtColor(roi_cv, cv2.COLOR_BGR2GRAY)

    # 检查是否与文本区域重叠过多
    region_area = (x2 - x1) * (y2 - y1)
    text_overlap = 0

    for text_bbox in text_regions:
        tx1, ty1, tx2, ty2 = text_bbox

        # 计算重叠面积
        overlap_x1 = max(x1, tx1)
        overlap_y1 = max(y1, ty1)
        overlap_x2 = min(x2, tx2)
        overlap_y2 = min(y2, ty2)

        if overlap_x1 < overlap_x2 and overlap_y1 < overlap_y2:
            text_overlap += (overlap_x2 - overlap_x1) * (overlap_y2 - overlap_y1)

    text_overlap_ratio = text_overlap / region_area

    # 如果文本重叠过多，可能不是纯图表
    if text_overlap_ratio > 0.7:
        return False

    # 检查几何形状特征
    edges = cv2.Canny(gray_roi, 50, 150)

    # 查找直线（架构图通常包含很多直线）
    lines = cv2.HoughLinesP(
        edges, 1, np.pi / 180, threshold=20, minLineLength=30, maxLineGap=10
    )
    line_count = len(lines) if lines is not None else 0

    # 查找圆形/椭圆（流程图节点）
    circles = cv2.HoughCircles(
        gray_roi,
        cv2.HOUGH_GRADIENT,
        1,
        20,
        param1=50,
        param2=30,
        minRadius=10,
        maxRadius=100,
    )
    circle_count = len(circles[0]) if circles is not None else 0

    # 查找矩形（架构图组件）
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    rect_count = 0
    for contour in contours:
        # 近似轮廓为多边形
        epsilon = 0.02 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)
        if len(approx) == 4:  # 矩形
            rect_count += 1

    # 评分系统
    score = 0
    if line_count > 5:
        score += 2
    if circle_count > 0:
        score += 2
    if rect_count > 2:
        score += 3
    if text_overlap_ratio < 0.3:
        score += 1

    return score >= 3


class DiagramExtractor:

    def __init__(self, min_diagram_size=(200, 200), dpi=300, frame_color=None):
        """
        架构图智能提取器

        Args:
            min_diagram_size: 最小图表尺寸 (width, height)
            dpi: 渲染分辨率
        """
        self.min_width, self.min_height = min_diagram_size
        self.dpi = dpi
        # 用于确保文档在需要时保持打开状态
        self.current_doc = None
        self.lock = threading.Lock()
        self.frame_color = frame_color  # 新增配置

    def _is_likely_diagram(self, region, page_image, text_regions):
        """
        判断区域是否可能是架构图。
        返回:
            (bool, bool): (is_diagram, was_identified_by_frame)
        """
        x1, y1, x2, y2 = region
        roi = page_image.crop((x1, y1, x2, y2))
        roi_cv = cv2.cvtColor(np.array(roi), cv2.COLOR_RGB2BGR)

        # 检查是否有指定颜色的边框
        if self.frame_color:
            if _has_colored_frame(roi_cv, color_name=self.frame_color):
                # 是图表，且是因为边框识别的
                return True, True

        # --- 后续逻辑不变，但返回值格式需要统一 ---

        # 检查是否与文本区域重叠过多
        region_area = (x2 - x1) * (y2 - y1)
        if region_area == 0:
            return False, False  # 不是图表
        text_overlap = 0
        # ... (计算 text_overlap 的代码不变)
        for text_bbox in text_regions:
            tx1, ty1, tx2, ty2 = text_bbox
            overlap_x1 = max(x1, tx1)
            overlap_y1 = max(y1, ty1)
            overlap_x2 = min(x2, tx2)
            overlap_y2 = min(y2, ty2)
            if overlap_x1 < overlap_x2 and overlap_y1 < overlap_y2:
                text_overlap += (overlap_x2 - overlap_x1) * (overlap_y2 - overlap_y1)

        text_overlap_ratio = text_overlap / region_area
        if text_overlap_ratio > 0.7:
            return False, False  # 不是图表

        # 检查几何形状特征...
        # ... (评分系统代码不变)
        gray_roi = cv2.cvtColor(roi_cv, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray_roi, 50, 150)
        lines = cv2.HoughLinesP(
            edges, 1, np.pi / 180, threshold=20, minLineLength=30, maxLineGap=10
        )
        line_count = len(lines) if lines is not None else 0
        circles = cv2.HoughCircles(
            gray_roi,
            cv2.HOUGH_GRADIENT,
            1,
            20,
            param1=50,
            param2=30,
            minRadius=10,
            maxRadius=100,
        )
        circle_count = len(circles[0]) if circles is not None else 0
        contours, _ = cv2.findContours(
            edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        rect_count = 0
        for contour in contours:
            epsilon = 0.02 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            if len(approx) == 4:
                rect_count += 1

        score = 0
        if line_count > 5:
            score += 2
        if circle_count > 0:
            score += 2
        if rect_count > 2:
            score += 3
        if text_overlap_ratio < 0.3:
            score += 1

        # 根据分数判断，但它不是通过边框识别的
        is_diagram = score >= 3
        return is_diagram, False

    def detect_diagram_regions(self, page_image):
        """
        检测页面中的图表区域

        Args:
            page_image: PIL Image对象

        Returns:
            list: 图表区域坐标列表 [(x1,y1,x2,y2), ...]
        """
        # 转换为OpenCV格式
        cv_image = cv2.cvtColor(np.array(page_image), cv2.COLOR_RGB2BGR)
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)

        # 边缘检测
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)

        # 形态学操作，连接边缘
        kernel = np.ones((3, 3), np.uint8)
        edges = cv2.dilate(edges, kernel, iterations=2)
        edges = cv2.erode(edges, kernel, iterations=1)

        # 查找轮廓
        contours, _ = cv2.findContours(
            edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        diagram_regions = []

        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)

            # 过滤太小的区域
            if w < self.min_width or h < self.min_height:
                continue

            # 过滤长宽比异常的区域（可能是文本行）
            aspect_ratio = w / h
            if aspect_ratio > 10 or aspect_ratio < 0.1:
                continue

            # 检查区域密度（边缘密度高的可能是图表）
            roi_edges = edges[y: y + h, x: x + w]
            edge_density = np.sum(roi_edges > 0) / (w * h)

            if edge_density > 0.01:  # 边缘密度阈值
                diagram_regions.append((x, y, x + w, y + h))

        return _merge_overlapping_regions(diagram_regions)

    def detect_vector_graphics(self, page):
        """
        检测PDF页面中的矢量图形区域
        """
        # 获取页面中的绘制指令
        drawings = page.get_drawings()

        if not drawings:
            return []

        # 分析绘制区域
        regions = []
        for drawing in drawings:
            if "rect" in drawing:
                rect = drawing["rect"]
                x1, y1, x2, y2 = rect.x0, rect.y0, rect.x1, rect.y1

                # 转换为像素坐标
                mat = fitz.Matrix(self.dpi / 72, self.dpi / 72)
                p1 = fitz.Point(x1, y1) * mat
                p2 = fitz.Point(x2, y2) * mat

                w = abs(p2.x - p1.x)
                h = abs(p2.y - p1.y)

                if w > self.min_width and h > self.min_height:
                    regions.append(
                        (
                            int(min(p1.x, p2.x)),
                            int(min(p1.y, p2.y)),
                            int(max(p1.x, p2.x)),
                            int(max(p1.y, p2.y)),
                        )
                    )

        return regions

    def extract_diagrams_from_pdf(self, pdf_path, base_output_dir):
        with self.lock:
            output_dir = base_output_dir
            pictures_dir = os.path.join(output_dir, "pictures")
            os.makedirs(pictures_dir, exist_ok=True)

            # 使用上下文管理器确保文档正确关闭
            with fitz.open(pdf_path) as doc:
                self.current_doc = doc  # 存储当前文档引用

                # *** 新增：初始化一个空列表来存储带框图表的名称 ***
                framed_diagram_names = []

                markdown_content = (
                    f"# {os.path.basename(pdf_path)} - 架构图提取结果\n\n"
                )
                diagram_count = 0

                # print(f"开始智能分析PDF: {pdf_path}")
                # print(f"总页数: {doc.page_count}")
                # print(f"输出目录: {output_dir}")

                for page_num in range(doc.page_count):
                    page = doc.load_page(page_num)

                    # 渲染页面为图像
                    mat = fitz.Matrix(self.dpi / 72, self.dpi / 72)
                    pix = page.get_pixmap(matrix=mat)  # type: ignore
                    page_image = Image.frombytes(
                        "RGB", (pix.width, pix.height), pix.samples
                    )

                    # 获取文本区域
                    text_regions = extract_text_blocks(page)

                    # 方法1：基于图像处理的区域检测
                    image_regions = self.detect_diagram_regions(page_image)

                    # 方法2：基于PDF矢量图形的检测
                    vector_regions = self.detect_vector_graphics(page)

                    # 合并所有候选区域
                    all_regions = image_regions + vector_regions
                    all_regions = _merge_overlapping_regions(all_regions)

                    # 提取文本内容
                    page_text = page.get_text().strip()  # type: ignore

                    if page_text or all_regions:
                        markdown_content += f"## 第 {page_num + 1} 页\n\n"

                        if page_text:
                            markdown_content += f"### 📝 文本内容\n\n{page_text}\n\n"

                    # 分析每个区域
                    page_diagrams = 0
                    for i, region in enumerate(all_regions):
                        # *** 这里是关键改动 ***
                        # 调用的是 self._is_likely_diagram
                        is_diagram, by_frame = self._is_likely_diagram(
                            region, page_image, text_regions
                        )

                        if is_diagram:
                            diagram_count += 1
                            page_diagrams += 1
                            # 提取图表区域
                            x1, y1, x2, y2 = region
                            diagram_img = page_image.crop((x1, y1, x2, y2))

                            # 保存图表
                            diagram_name = f"diagram_{diagram_count:03d}_page{page_num + 1}_{i + 1}.png"
                            diagram_path = os.path.join(pictures_dir, diagram_name)
                            diagram_img.save(diagram_path, "PNG", quality=95)

                            # 添加到Markdown
                            markdown_content += f"![](./pictures/{diagram_name})\n\n"
                            # markdown_content += f"### 🏗️ 架构图 {page_diagrams}\n\n"
                            # markdown_content += f"![架构图 - 第{page_num + 1}页](./pictures/{diagram_name})\n\n"
                            # markdown_content += f"*位置: 第{page_num + 1}页，区域{i + 1}*\n\n

                            if by_frame:
                                framed_diagram_names.append(diagram_name)

                # 保存结果
                file_name_no_ext = os.path.splitext(os.path.basename(pdf_path))[0]
                md_path = os.path.join(output_dir, f"{file_name_no_ext}.md")
                with open(md_path, "w", encoding="utf-8") as f:
                    f.write(markdown_content)

                # print(f"\n🎉 提取完成！")
                # print(f"📊 总共提取了 {diagram_count} 个架构图")
                # print(f"📄 主文档: {md_path}")
                # print(f"🖼️  架构图目录: {pictures_dir}")

                self.current_doc = None  # 处理完成后清除引用
                return md_path, diagram_count, framed_diagram_names


async def insert_image_descriptions_to_md(
        md_path, base_output_dir, framed_diagram_names, language="中文"
):
    """
    对指定图片列表，调用 LLM 获取解释，并插入到 markdown 图片下方
    """
    image_desc_map = {}
    if framed_diagram_names:
        from pathlib import Path
        pictures_dir = os.path.join(base_output_dir, "pictures")
        # 串行调用 LLM
        desc_list = []
        for img_name in framed_diagram_names:
            img_path = Path(pictures_dir) / img_name
            try:
                desc = await describe_image_with_llm(img_path, language=language)
                desc_list.append(desc)
            except Exception as e:
                # Log the error but continue with other images
                print(f"Failed to describe image {img_name}: {str(e)}")
                desc_list.append("")  # Add empty description for failed image
        image_desc_map = dict(zip(framed_diagram_names, desc_list))
    if not image_desc_map:
        return
    with open(md_path, "r", encoding="utf-8") as f:
        md_content = f.read()

        def insert_desc(md):
            for img_name, desc in image_desc_map.items():
                # Only insert description if it's not empty
                if desc.strip():
                    pattern = f"![](./pictures/{img_name})"
                    replacement = f"{pattern}\n> 图片解释：{desc.strip()}\n"
                    md = md.replace(pattern, replacement)
            return md

        md_content = insert_desc(md_content)
    with open(md_path, "w", encoding="utf-8") as f:
        f.write(md_content)


async def pdf_convert_to_md(file_path: Path, language: str):
    """
    智能转换 PDF 为 Markdown
    """
    # 1. 文档转换
    base_output_dir = str(file_path.parent)
    frame_color = getattr(settings, "pdf_picture_need_desc_frame_color", None)
    extractor = DiagramExtractor(
        min_diagram_size=(150, 150), dpi=300, frame_color=frame_color
    )
    md_path, _, framed_diagram_names = extractor.extract_diagrams_from_pdf(
        str(file_path), base_output_dir
    )
    # 2. 插入图片解释
    await insert_image_descriptions_to_md(
        md_path, base_output_dir, framed_diagram_names, language
    )
    # 3. 移动目录到 nginx_static_dir 下, 替换 md 文件中的图片路径
    nginx_static_dir = settings.nginx_static_dir
    nginx_url_prefix = settings.nginx_url_prefix
    dest_nginx_dir = os.path.join(nginx_static_dir, os.path.basename(base_output_dir))
    shutil.copytree(base_output_dir, dest_nginx_dir, dirs_exist_ok=True)
    with open(md_path, "r", encoding="utf-8") as f:
        md_content = f.read()
        md_content = md_content.replace(
            "./pictures/",
            f"{nginx_url_prefix}/{os.path.basename(base_output_dir)}/pictures/",
        )
    with open(md_path, "w", encoding="utf-8") as f:
        f.write(md_content)
    return Path(md_path)
